import os
import sys
from transformers import pipeline, ViTImageProcessor, ViTForImageClassification

# 模型和图像路径
model_path = "/opt/model/google/vit-base-patch16-224"
image_path = "./data/image/pipeline-cat-chonk.jpeg"

def verify_pillow_installation():
    """强制验证Pillow是否能正常导入"""
    try:
        import PIL
        from PIL import Image
        print(f"Pillow已安装，版本：{PIL.__version__}")
        return True
    except ImportError:
        raise ImportError(
            "Pillow库未正确安装！\n"
            "请执行以下命令重新安装：\n"
            "pip3 uninstall -y Pillow\n"
            "pip3 install Pillow --no-cache-dir"
        )

try:
    # 首先验证Pillow
    verify_pillow_installation()
    
    # 验证模型路径
    if not os.path.exists(model_path):
        raise FileNotFoundError(f"模型路径不存在: {model_path}")
    
    # 验证图像文件
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"图像文件不存在: {image_path}")
    
    # 加载处理器和模型（这一步需要Pillow）
    print("加载图像处理器...")
    processor = ViTImageProcessor.from_pretrained(model_path, local_files_only=True)
    
    print("加载模型...")
    model = ViTForImageClassification.from_pretrained(model_path, local_files_only=True)
    
    # 图像分类
    classifier = pipeline(
        "image-classification",
        model=model,
        image_processor=processor
    )
    
    result = classifier(image_path)
    print("图像分类结果：", result)

except Exception as e:
    print(f"执行失败: {str(e)}", file=sys.stderr)
